# -*- coding: UTF-8 -*-
import os
import Queue
import socket
import sys
import threading
import time
import urllib
from urlparse import urlparse

import requests

socket.setdefaulttimeout(60)  # 网络超时60s


def get_url_list(host, body):
    urlgroup = urlparse(host)
    hostWithoutPath = urlgroup.scheme + '://' + urlgroup.hostname
    lines = body.split('\n')
    ts_url_list = []
    i = 1
    for line in lines:
        if not line.startswith('#') and line != '':
            if line.endswith('m3u8'):
                if line.find('/') != -1:
                    url = combineUrl(hostWithoutPath, line)
                else:
                    url = '%s/%s' % (host, line)
                return downloadWrap(url)
            elif line.startswith('http'):
                url = line
            elif line.find('/') != -1:
                url = combineUrl(hostWithoutPath, line)
            else:
                url = '%s/%s' % (host, line)
            # status -1 未开始下载 0 下载中 1 已完成
            ts_url_list.append(
                {'url': url, 'status': -1, 'index': i, 'ts_name': line})
            i += 1
    return ts_url_list

# 拼接链接


def combineUrl(baseurl, relateurl):
    if relateurl.startswith('/'):
        url = '%s%s' % (baseurl, relateurl)
    else:
        url = '%s/%s' % (baseurl, relateurl)
    return url


def get_host(url):
    urlgroup = urlparse(url)
    if url.index('.m3u8'):
        return url[:url.rindex('/')]
    return urlgroup.scheme + '://' + urlgroup.hostname


def get_m3u8_body(url):
    print 'read m3u8 file:', url
    session = requests.Session()
    adapter = requests.adapters.HTTPAdapter(
        pool_connections=10, pool_maxsize=10, max_retries=10)
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    r = session.get(url, timeout=10)
    return r.content


# 下载单个ts文件
def dowload_ts(thread, tsItem):
    global download_count
    lock.acquire()
    if not os.path.isdir(temp_dir):
        os.makedirs(temp_dir)
    lock.release()
    ts_url = tsItem['url']
    ts_name = ts_url[ts_url.rfind('/'):]
    curr_path = '%s%s' % (temp_dir, ts_name)
    try:
        if os.path.isfile(curr_path):
            print '\n'+curr_path+'文件已存在！'
        else:
            filename, h = urllib.urlretrieve(ts_url, curr_path)
            if int(h['Content-Length']) < 1000:
                print ts_url+' 链接下载异常，请检查网络'
                os.remove(curr_path)
                os._exit(0)
        lock.acquire()
        tsItem['status'] = 1
        download_count += 1
        lock.release()
    except:
        lock.acquire()
        tsItem['status'] = -1
        if os.path.isfile(curr_path):
            os.remove(curr_path)
        lock.release()
    lock.acquire()
    toolBar(download_count)
    lock.release()
    thread.run()


def toolBar(num):
    global precent
    has_loaded_list = filter(lambda item: item['status'] == 1, ts_url_list)
    has_loaded_count = len(has_loaded_list)
    new_precent = int(has_loaded_count*100/total_count)
    if precent < new_precent:
        precent = new_precent
        sys.stdout.write('\r')
        sys.stdout.write("[%s>%s] %s" %
                         ('-'*precent, ' '*(100-precent), str(precent)+'%'))
        sys.stdout.flush()
        print ""


# 合并文件
def combine_file(filename):
    # 合并ts文件
    if not os.path.isdir(download_dir):
        os.makedirs(download_dir)
    os.chdir(temp_dir)
    shell_str = ' '.join(map(lambda item: (item['ts_name'], item['ts_name'][item['ts_name'].rfind(
        '/'):].replace('/', ''))[item['ts_name'].rfind('/') != -1], ts_url_list))
    try:
        shell_str = 'cat ' + shell_str + ' > ' + \
            download_dir + '/' + filename
        os.system(shell_str)
        print '合并完成 生成视频 ' + filename
        # 删除ts和m3u8文件
        os.system('rm -f *.ts')
        os.system('rm -f *.m3u8')
        threadsEnd()
    except:
        os._exit(0)


def main():
    global temp_dir, download_dir, total_count, download_count, ts_url_list, workQueue, precent, filename
    workQueue = Queue.Queue(1)
    url = ''
    while not url:
        url = raw_input('请输入m3u8文件路径：')
        # if not url.endswith('.m3u8'):
        #     url = ''
        #     print '文件后缀不是m3u8格式'
        # else:
        #     if (url.startswith('http://') == False) and (url.startswith('https://') == False):
        #         url = ''
        #         print '输入的不是合法的网络路径'

    # 生成的视频文件存放的文件夹
    download_dir = raw_input('请输入视频生成的文件夹[./video]：')
    if download_dir:
        while len(download_dir) and os.path.isdir(download_dir) == False:
            download_dir = raw_input('请输入视频生成的文件夹[./video]：')
    else:
        download_dir = './video'
    download_dir = os.path.abspath(download_dir)
    filename = raw_input('生成的视频名称[video1.mp4]：')
    if filename:
        filename += '.mp4'
    else:
        filename += 'video1.mp4'
    # 下载的ts文件存储的文件夹
    temp_dir = './.temp'
    # 总下载文件数
    total_count = 0
    # 已经下载的文件数
    download_count = 0

    precent = 0  # 当前下载进度
    downloadWrap(url)


def downloadWrap(url):
    global ts_url_list, lock, total_count
    host = get_host(url)
    body = get_m3u8_body(url)
    ts_url_list = get_url_list(host, body)
    if ts_url_list == None or not len(ts_url_list):
        return
    total_count = len(ts_url_list)
    unload_count = total_count
    lock = threading.Lock()
    count = 0
    while count < 5:
        generatorThreads(unload_count)
        threadsEnd()
        unload_count = len(
            filter(lambda item: item['status'] != 1, ts_url_list))
        if unload_count:
            print '%d个文件未下载完成，重新下载' % unload_count
            count += 1
        else:
            break
    if not unload_count:
        combine_file(filename)
    else:
        print '部分文件下载失败!!!'


# 开启多线程
def generatorThreads(unload_count):
    global threadList
    threadCount = 20
    if unload_count < threadCount:
        threadCount = unload_count
    threadList = []
    for i in range(threadCount):
        thread = myThread(i, 'thread-'+str(i))
        thread.start()
        threadList.append(thread)


# 线程结束
def threadsEnd():
    for thread in threadList:
        thread.join()

# 自定义线程操作


class myThread(threading.Thread):
    def __init__(self, threadID, name):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = name

    def run(self):
        lock.acquire(1)
        unloadItem = getUnloadItem(ts_url_list)
        lock.release()
        if unloadItem:
            dowload_ts(self, unloadItem)

# 获取未下载的ts链接


def getUnloadItem(alist):
    for item in alist:
        if item != None and item['status'] == -1:
            item['status'] = 0
            return item
        else:
            continue
    return None


if __name__ == '__main__':
    main()
